from operator import itemgetter
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.runnables import RunnableWithMessageHistory

from biz.core.sql.history import LimitedSQLChatMessageHistory
from biz.config.settings import llm
from biz.core.ai.tools import road_function, pole_function, ask_for_more_info

parser = JsonOutputParser()
def get_session_history(session_id):
    connection = "sqlite:///memory.db"
    sql_chat_message_history = LimitedSQLChatMessageHistory(session_id, connection)
    return sql_chat_message_history

def tool_chain(model_output):
    tool_map = {tool.name: tool for tool in [road_function, pole_function, ask_for_more_info]}
    chosen_tool = tool_map[model_output["name"]]
    return itemgetter("arguments") | chosen_tool

def get_args(res):
    args=res.split("Action:")[1].split("Action Input:")
    try:
        arguments = parser.parse(args[1].strip())
    except:
        arguments = {}
    return {"name": args[0].strip(), "arguments": arguments}


runnable_with_history = RunnableWithMessageHistory(
    llm,
    get_session_history,
)

def conversation_context_builder(response, conversation_context):
    lines = response.split('\n')
    # 遍历每一行并提取关键信息
    for line in lines:
        if line.startswith('Thought:'):
            conversation_context = conversation_context + "\n" + line
        elif line.startswith('Action:'):
            conversation_context = conversation_context + "\n" + line
        elif line.startswith('Action Input:'):
            conversation_context = conversation_context + "\n" + line
    return conversation_context


def get_response(response, conversation_context):
    response = response.replace("Observation", "")
    final_answer = ''
    # 返回的信息先进行处理
    lines = response.split('\n')
    func_context = ''
    func_list = []
    for line in lines:
        if line.startswith('Thought:'):
            func_context = func_context + "\n" + line
        elif line.startswith('Action:'):
            func_context = func_context + "\n" + line
        elif line.startswith('Action Input:'):
            func_context = func_context + "\n" + line
        elif line.startswith('Final Answer:'):
            final_answer = line.replace("Final Answer:", "")

        if func_context.find("Action:") != -1 and func_context.find("Action Input:") != -1:
            func_list.append(func_context)
            func_context = ''

    if len(func_list) == 0:
        if final_answer == '':
            return False, None, '回复格式不正确，请按正确格式回复!!!'
        return True, final_answer, conversation_context

    result = []
    for func in func_list:
        success, temp_res, conversation_context = single_method(conversation_context, func)
        result.append(temp_res)
        if not success:
            return False, '<br>'.join(str(item) for item in result), conversation_context
    return success, '<br>'.join(str(item) for item in result), conversation_context


def single_method(conversation_context, response):
    if response.rfind("Action Input:") != -1:
        try:
            args = get_args(response)
            t_run = tool_chain(args)
            if args["name"] == "ask_for_more_info":
                question = args["arguments"]["question"]
                print("\033[1;31;40m" + question + "\033[0m")
                return False, question, conversation_context
            else:
                success, query = t_run.invoke(args)
                if not success:
                    return success, None, query
                return success, query, conversation_context
        except Exception as e:
            print(e)
            return True, "未找到可以调用的工具!!!", conversation_context
    elif response.rfind("Final Answer:") != -1:
        final_answer = response.split("Final Answer:")[1].strip()
        print("\033[1;31;40m" + final_answer + "\033[0m")
        return True, final_answer, conversation_context
    else:
        query = "回复格式不正确，请按正确格式回复!!!"
        conversation_context = conversation_context + "\nObservation: " + query
        return False, None, conversation_context